#coding:utf-8
'''
预测输出结果
1. 预测数据读取+预处理
2. 应用模型, 得到结果
3. 保存结果
'''
import torch
import numpy as np
import csv
from datasets import load_from_disk
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

'''
1. 预测数据读取+预处理
'''
#从本地磁盘加载数据
datasets = load_from_disk('./data/ChnSentiCorp')
#取数据子集，否则数据太多跑不动
dataset_test = datasets['test'].shuffle().select(range(100))    #200
del datasets
'''
分词处理
'''
#加载分词工具
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')   #---选定分词器
#针对每个数据进行分词
def f(data):
    return tokenizer(
        data['text'],         #---查看前面数据的内容
        padding='max_length', #一律补零到max_length长度
        truncation=True,      #超过就截断
        max_length=30,
    )
#批量处理
dataset_test = dataset_test.map(f, batched=True, batch_size=100, num_proc=4)
'''
DataLoader
'''
# 将默认数据格式list转换为tensor
def collate_fn(data):
    label = [i['label'] for i in data]
    input_ids = [i['input_ids'] for i in data]
    token_type_ids = [i['token_type_ids'] for i in data]
    attention_mask = [i['attention_mask'] for i in data]

    # 数据类型格式的转换
    label = torch.LongTensor(label)
    input_ids = torch.LongTensor(input_ids)
    token_type_ids = torch.LongTensor(token_type_ids)
    attention_mask = torch.LongTensor(attention_mask)

    return label, input_ids, token_type_ids, attention_mask

# 加载数据集
loader_test = torch.utils.data.DataLoader(dataset=dataset_test,     #----改变数据集
                                          batch_size=4,
                                          collate_fn=collate_fn,    # 数据格式的处理
                                          shuffle=True,
                                          drop_last=True)

'''
2. 应用模型, 得到结果
'''
#加载模型
model = AutoModelForSequenceClassification.from_pretrained('bert-base-cased',  #---模型
                                                           num_labels=2)       #---类别
def test(loader_test):
    #加载参数：保存模型后
    model.load_state_dict(torch.load('./output_dir/pytorch_model.bin'))
    
    outs = np.array([],dtype=int)    #预测值
    #labels = np.array([],dtype=int)  #真实值
    
    for i, (label, input_ids, token_type_ids,attention_mask) in enumerate(loader_test):
        # 查看前5个数据预测效果
        #if i == 5:
        #    break
        
        # 不进行梯度下降
        with torch.no_grad():
            out = model(input_ids=input_ids,
                        attention_mask=attention_mask,
                        token_type_ids=token_type_ids)
        #[4, 2] -> [4]
        out = out['logits'].argmax(dim=1)
        
        outs = np.append(outs, out)
        #labels = np.append(labels, label)

    return {'pre':outs, 'index':i}
    #return {'pre':outs, 'label':labels}

saveDict = test(loader_test)
print(saveDict)

'''
3. 保存结果【有问题，待解决】

fileName="./output_dir/predict.csv"
##保存文件
with open(fileName,"wb") as csv_file:
	writer=csv.writer(csv_file)
	for key,value in saveDict.items:
		writer.writerow([key,value])
        '''